#!/usr/bin/env python3
import rospy
import yaml
import os
from geometry_msgs.msg import PoseWithCovarianceStamped, PoseStamped

# 全局变量来存储最新位姿
current_pose = None
pose_topic = "/amcl_pose"

def pose_callback(data):
    """回调函数，每次收到位姿更新时，存储最新位姿"""
    global current_pose
    if isinstance(data, PoseWithCovarianceStamped):
        current_pose = PoseStamped()
        current_pose.header = data.header
        current_pose.pose = data.pose.pose
    elif isinstance(data, PoseStamped):
        current_pose = data
    else:
        rospy.logwarn("收到的位姿消息格式不正确。")

def save_goal_to_yaml(file_path, goal_name, pose):
    """将新的目标点保存到YAML文件中"""
    # 检查文件是否存在，如果不存在则创建
    if not os.path.exists(file_path):
        data = {'goals': {}}
    else:
        # 读取现有文件内容
        try:
            with open(file_path, 'r') as f:
                data = yaml.safe_load(f) or {'goals': {}}
        except yaml.YAMLError as e:
            rospy.logerr(f"读取 YAML 文件失败: {e}")
            return False

    # 准备新目标点的数据
    new_goal_data = {
        'position': [
            round(pose.position.x, 3),
            round(pose.position.y, 3),
            round(pose.position.z, 3)
        ],
        'orientation': [
            round(pose.orientation.x, 3),
            round(pose.orientation.y, 3),
            round(pose.orientation.z, 3),
            round(pose.orientation.w, 3)
        ]
    }

    # 将新目标点添加到字典中
    data['goals'][goal_name] = new_goal_data

    # 将更新后的数据写入文件
    try:
        with open(file_path, 'w') as f:
            yaml.dump(data, f, default_flow_style=None)
            rospy.loginfo(f"已成功将目标点 '{goal_name}' 保存到: {file_path}")
            return True
    except IOError as e:
        rospy.logerr(f"写入文件失败: {e}")
        return False

def main():
    global current_pose, pose_topic
    rospy.init_node('goal_file_manager', anonymous=False)
    
    # 从 launch 文件中获取 YAML 文件路径
    goal_file_path = rospy.get_param('~goal_file_path', 'goals.yaml')
    
    # 从 launch 文件中获取位姿话题名称
    pose_topic = rospy.get_param('~pose_topic', '/amcl_pose')

    rospy.Subscriber(pose_topic, PoseWithCovarianceStamped, pose_callback)

    rospy.loginfo(f"位姿记录节点已启动，监听话题: {pose_topic}")
    rospy.loginfo(f"目标点将保存到: {goal_file_path}")
    rospy.loginfo("使用遥控器移动机器人到目标位置，然后输入目标点名称并回车。")

    while not rospy.is_shutdown():
        # 等待用户输入目标点名称
        goal_name = input("请输入目标点名称 (例如: kitchen): ")
        
        if not goal_name:
            rospy.logwarn("目标点名称不能为空。")
            continue
        
        if current_pose is None:
            rospy.logwarn("无法获取当前位姿。请确保机器人位姿话题正在发布数据。")
            continue

        # 保存位姿到文件
        if save_goal_to_yaml(goal_file_path, goal_name, current_pose.pose):
            rospy.loginfo(f"目标点 '{goal_name}' 已成功保存。")
        else:
            rospy.logerr(f"保存目标点 '{goal_name}' 失败。")

if __name__ == '__main__':
    try:
        main()
    except rospy.ROSInterruptException:
        pass